from bokeh.io import curdoc
from bokeh.layouts import row, column
from bokeh.models import (ColumnDataSource, PointDrawTool, 
                         Div, Toggle, TextInput, Paragraph,
                         RadioGroup, Arrow, OpenHead)  # 新增箭头相关组件
from bokeh.models import RadioGroup
from bokeh.plotting import figure
import numpy as np
from utils import generate_colors, validate_coordinates, edges_to_bokeh, create_editable_table, create_readonly_table
from metric_functions import custom_metric, calculate_centroid, adjacency_edges, geometry_distance
from functools import partial
import torch

CENTROID_COLOR = "#2C3E50"  # 中性深灰色

class VisualizationApp:
    def __init__(self, point_count=4, k_neighbors=2):
        # === 初始化参数 ===
        self.point_count = point_count
        self.k = k_neighbors
        self.show_edges = False  # 邻接边显示状态
        
        # === 数据源初始化 ===
        colors = generate_colors(self.point_count)
        
         # 主点数据（新增color_hex字段用于表格颜色显示）
        self.source_left = ColumnDataSource({
            'x': np.random.rand(point_count)*10,
            'y': np.random.rand(point_count)*10,
            'color': colors,
            'color_hex': colors  # 新增颜色编码字段
        })
        self.source_right = ColumnDataSource({
            'x': np.random.rand(point_count)*10,
            'y': np.random.rand(point_count)*10,
            'color': colors,
            'color_hex': colors
        })
        
         # 修复质心数据源（初始长度一致）
        self.centroid_left = ColumnDataSource({'x': [], 'y': [], 'color_hex': []})
        self.centroid_right = ColumnDataSource({'x': [], 'y': [], 'color_hex': []})
        
        # 边数据源（存储邻接边坐标）
        self.edge_source_left = ColumnDataSource({'xs': [], 'ys': []})
        self.edge_source_right = ColumnDataSource({'xs': [], 'ys': []})

        # === 新增梯度相关属性 ===
        self.selected_gradient = "none"  # 初始状态对应"不显示梯度"
        self.grad_arrows = {'left': None, 'right': None}  # 梯度箭头渲染器
        
        # === 新增控件 ===
        self.k_input = TextInput(
            title="邻接边数量k:", 
            value=str(k_neighbors), 
            width=100,
            styles={'margin-right': '10px'}
        )

        # === 新增控件 ===
        self.lam_geo_input = TextInput(
            title="λ_geodesic:", 
            value="1.0",
            width=80,
            styles={'margin-right': '10px'}
        )
        self.lam_geom_input = TextInput(
            title="λ_geometry:", 
            value="0.0",
            width=80,
            styles={'margin-right': '10px'}
        )
        
        self.gradient_selector = RadioGroup(
            labels=["不显示梯度", "Geodesic梯度", "Geometry梯度", "HT2T梯度"],  # 新增首选项
            active=0,  # 默认不显示
            width=180
        )

        # 新增点配置控件
        self.point_config_selector = RadioGroup(
            labels=["预设点", "随机点"],
            active=0,
            width=120
        )

        # 新增部分函数回调绑定
        self.left_callback = partial(self._safe_data_update, 'left')
        self.right_callback = partial(self._safe_data_update, 'right')

        # 新增事件绑定
        self.lam_geo_input.on_change('value', self._update_lambda)
        self.lam_geom_input.on_change('value', self._update_lambda)
        self.gradient_selector.on_change('active', self._update_gradient_display)
        self.point_config_selector.on_change('active', self._update_point_config)
        
        # === 界面组件初始化 ===
        self._create_plots()      # 创建坐标系
        self._create_widgets()   # 创建控制部件
        self._add_drag_tools()    # 添加拖拽工具
        self._create_tables()     # 创建坐标表格
        
        # === 事件绑定 ===
        self._bind_events()
        
        # === 初始渲染 ===
        self._update_display()

        # 新增初始化清除
        self._plot_gradients(np.zeros((point_count,2)), 
                        np.zeros((point_count,2)))

        

    def _create_plots(self):
        """创建左右坐标系（修复diamond方法弃用问题）"""
        # 左图配置
        self.plot_left = figure(
            width=400, 
            height=400,
            tools="",  # 延迟添加工具避免冲突
            title="左坐标系",
            toolbar_location="right"
        )
        # 主点集（可拖拽）
        self.main_renderer_left = self.plot_left.scatter(
            'x', 'y', 
            source=self.source_left,
            size=15, 
            color='color', 
            alpha=0.8
        )
        # 中点（使用scatter替代diamond方法）
        self.plot_left.scatter(
            'x', 'y', 
            source=self.centroid_left,
            size=20,
            marker="diamond",  # 指定菱形标记
            fill_color=CENTROID_COLOR,
            line_color="black",
            line_width=1.5,
            alpha=0.8
        )
        # 邻接边（初始隐藏）
        self.plot_left.multi_line(
            'xs', 'ys',
            source=self.edge_source_left,
            line_color="#E74C3C", 
            line_alpha=0.6
        )
        self.plot_left.add_tools('pan','wheel_zoom','reset')

        # 右图配置（结构与左图对称）
        self.plot_right = figure(
            width=400,
            height=400,
            tools="",
            title="右坐标系",
            toolbar_location="right"
        )
        self.main_renderer_right = self.plot_right.scatter(
            'x', 'y', 
            source=self.source_right,
            size=15,
            color='color', 
            alpha=0.8
        )
        self.plot_right.scatter(
            'x', 'y',
            source=self.centroid_right,
            size=20,
            marker="diamond",
            fill_color=CENTROID_COLOR,
            line_color="black",
            line_width=1.5,
            alpha=0.8
        )
        self.plot_right.multi_line(
            'xs', 'ys',
            source=self.edge_source_right,
            line_color="#E74C3C",
            line_alpha=0.6
        )
        self.plot_right.add_tools('pan','wheel_zoom','reset')

    def _create_tables(self):
        """创建坐标表格（新增质心坐标表）"""
        self.editable_table_left = create_editable_table(self.source_left, "左图点坐标（可编辑）")
        self.centroid_table_left = create_readonly_table(self.centroid_left, "左图质心坐标")
        self.editable_table_right = create_editable_table(self.source_right, "右图点坐标（可编辑）")
        self.centroid_table_right = create_readonly_table(self.centroid_right, "右图质心坐标")

    def _create_widgets(self):
        """唯一控件创建方法（关键修复）"""
        # 主标题
        self.header = Div(
            text="<h1 style='text-align:center'>ht2t可视化（欧式版本）</h1>",
            styles={'margin-bottom': '20px'}
        )
        
        # 新增公式显示
        formula_html = """
        <div style="margin: 10px 0; font-size:12px">
            HT2T距离 = λ<sub>geodesic</sub> × 测地线距离 + λ<sub>geometry</sub> × 几何距离
        </div>
        """
        self.formula_display = Div(text=formula_html)

        # 邻接边切换按钮
        self.toggle = Toggle(
            label=f"显示邻接边 (k={self.k})", 
            active=False,
            button_type="primary",
            width=150
        )
        self.toggle.on_click(self._toggle_edges)
        
        # 差异显示组件修改
        self.geo_para = Paragraph(text="测地线距离: -", styles={'color': '#27AE60'})
        self.geom_para = Paragraph(text="几何距离: -", styles={'color': '#2980B9'})
        self.ht2t_para = Paragraph(text="HT2T距离: -", styles={'color': '#8E44AD', 'font-weight': 'bold'})
        self.grad_para = Paragraph(text="当前梯度: 未计算", styles={'color': '#E67E22'})

        # 修改布局
        self.diff_display = column(
            row(self.toggle, self.k_input, self.point_config_selector, align="center", spacing=20),
            row(self.lam_geo_input, self.lam_geom_input, self.formula_display),
            self.geo_para,
            self.geom_para,
            self.ht2t_para,
            row(self.gradient_selector, self.grad_para),
            styles={'margin-top': '20px'}
        )

        

    def _update_lambda(self, attr, old, new):
        """更新lambda参数"""
        try:
            self._update_metric()
        except:
            pass

    def _add_drag_tools(self):
        """添加点拖拽工具（适配Bokeh 3.4.1）"""
        for plot, renderer in [(self.plot_left, self.main_renderer_left),
                            (self.plot_right, self.main_renderer_right)]:
            tool = PointDrawTool(
                renderers=[renderer],
                description="拖拽点",
                empty_value='black',
                add=False,  # 禁用添加新点
                num_objects=self.point_count  # 固定点数量
            )
            plot.add_tools(tool)
            plot.toolbar.active_drag = tool

    def _bind_events(self):
        self.source_left.on_change('data', self.left_callback)
        self.source_right.on_change('data', self.right_callback)
        self.k_input.on_change('value', self._update_k_value)

    def _update_k_value(self, attr, old, new):
        """动态更新k值"""
        try:
            k = int(new)
            if 1 <= k < self.point_count:
                self.k = k
                self.toggle.label = f"显示邻接边 (k={k})"
                self._update_edges()
        except ValueError:
            pass

    def _safe_data_update(self, side, attr, old, new):
        """安全更新数据（强化校验）"""
        # 获取当前数据源
        source = self.source_left if side == 'left' else self.source_right
        callback = self.left_callback if side == 'left' else self.right_callback
        
        try:
            # 暂停事件监听
            source.remove_on_change('data', callback)
            
            # 强制数据维度对齐
            new_x = list(map(float, new['x']))[:self.point_count]
            new_y = list(map(float, new['y']))[:self.point_count]
            
            # 填充/截断数据
            new_x += [source.data['x'][i] for i in range(len(new_x), self.point_count)]
            new_y += [source.data['y'][i] for i in range(len(new_y), self.point_count)]
            
            # 执行安全更新
            source.patch({
                'x': [(slice(self.point_count), new_x)],
                'y': [(slice(self.point_count), new_y)]
            })
            
        except Exception as e:
            print(f"数据更新失败: {str(e)}")
            # 恢复原始数据
            source.data = dict(source.data)
        finally:
            # 重新绑定事件
            source.on_change('data', callback)
            self._update_display()

    def _toggle_edges(self, active):
        """切换邻接边显示状态"""
        self.show_edges = active
        self._update_edges()

    def _update_display(self):
        """统一更新入口（保证更新顺序）
        
        修改说明：
        - 新增梯度显示状态判断，确保数据变化时刷新梯度
        """
        self._update_centroids()  # 1. 更新中点
        self._update_edges()      # 2. 更新邻接边
        self._update_metric()     # 3. 更新差异值
        
        # 新增强制刷新逻辑
        if self.selected_gradient != "none":
            self._update_gradients()

    def _update_centroids(self):
        """修复质心计算（强制刷新数据源）"""
        for src, dst in [(self.source_left, self.centroid_left),
                    (self.source_right, self.centroid_right)]:
            # 强制转换为numpy数组
            points = np.column_stack([
                np.asarray(src.data['x'], dtype=np.float64),
                np.asarray(src.data['y'], dtype=np.float64)
            ])
            cx, cy = calculate_centroid(points)
            
            # 更新时重建数据字典
            dst.data = {
                'x': [float(cx)],
                'y': [float(cy)],
                'color_hex': [CENTROID_COLOR]
            }

    def _update_edges(self):
        """根据当前状态更新邻接边"""
        if self.show_edges:
            # 计算左图边
            left_points = np.column_stack([
                self.source_left.data['x'], 
                self.source_left.data['y']
            ])
            left_edges = adjacency_edges(left_points, self.k)
            self.edge_source_left.data = edges_to_bokeh(left_edges)
            
            # 计算右图边
            right_points = np.column_stack([
                self.source_right.data['x'], 
                self.source_right.data['y']
            ])
            right_edges = adjacency_edges(right_points, self.k)
            self.edge_source_right.data = edges_to_bokeh(right_edges)
        else:
            # 清空边数据
            self.edge_source_left.data = {'xs': [], 'ys': []}
            self.edge_source_right.data = {'xs': [], 'ys': []}

    def _update_metric(self):
        """更新双指标（修复版本）"""
        left_points = validate_coordinates(
            np.column_stack([
                self.source_left.data['x'], 
                self.source_left.data['y']
            ])
        )
        right_points = validate_coordinates(
            np.column_stack([
                self.source_right.data['x'], 
                self.source_right.data['y']
            ])
        )
        
        # 计算测地线距离
        centroid_left = calculate_centroid(left_points)
        centroid_right = calculate_centroid(right_points)
        geo_dist = np.linalg.norm(centroid_left - centroid_right)
        
        # 计算几何距离（添加异常处理）
        try:
            geom_dist = geometry_distance(left_points, right_points)
        except ValueError as e:
            print(f"几何距离计算错误: {str(e)}")
            geom_dist = float('nan')
        except Exception as e:
            print(f"几何距离计算错误: {str(e)}")
            geom_dist = float('nan')
        
        # 其他计算逻辑保持不变...
        lam_geo = float(self.lam_geo_input.value)
        lam_geom = float(self.lam_geom_input.value)
        ht2t_dist = lam_geo * geo_dist + lam_geom * (0 if np.isnan(geom_dist) else geom_dist)

        # 更新显示信息
        self.geo_para.text = f"测地线距离: {geo_dist:.4f}"
        self.geom_para.text = f"几何距离: {geom_dist:.4f}" if not np.isnan(geom_dist) else "几何距离: 计算错误"
        self.ht2t_para.text = f"HT2T距离: {ht2t_dist:.4f} (λ={lam_geo}, {lam_geom})"

        # 梯度更新逻辑保持不变...
        if self.selected_gradient != "none":
            self._update_gradients()
        else:
            self._plot_gradients(np.zeros_like(left_points), np.zeros_like(right_points))

    def _update_gradients(self):
        """计算并可视化梯度（显示分项梯度）"""
        if self.selected_gradient == "none":
            self.grad_para.text = "当前梯度: 未显示"
            return
        
        # 转换为可求导张量
        left_np = np.column_stack([self.source_left.data['x'], self.source_left.data['y']])
        right_np = np.column_stack([self.source_right.data['x'], self.source_right.data['y']])
        
        # 统一转换为浮点张量
        left_points = torch.tensor(left_np, dtype=torch.float32, requires_grad=True)
        right_points = torch.tensor(right_np, dtype=torch.float32, requires_grad=True)
        
        # 修改梯度计算分支
        if self.selected_gradient == "geodesic":
            loss = -custom_metric(left_points, right_points)
        elif self.selected_gradient == "geometry":
            loss = -geometry_distance(left_points, right_points)
        elif self.selected_gradient == "ht2t":
            geo = -custom_metric(left_points, right_points)
            geom = -geometry_distance(left_points, right_points)  # 保持张量计算
            loss = float(self.lam_geo_input.value)*geo + float(self.lam_geom_input.value)*geom

        # 反向传播
        if loss.requires_grad:
            loss.backward()
        
        # 获取梯度数据
        grads_left = left_points.grad.numpy() if left_points.grad is not None else np.zeros_like(left_np)
        grads_right = right_points.grad.numpy() if right_points.grad is not None else np.zeros_like(right_np)
        
        # 可视化梯度方向
        self._plot_gradients(grads_left, grads_right)
        
        # 计算分项梯度范数
        avg_grad_left = np.mean(np.linalg.norm(grads_left, axis=1))
        avg_grad_right = np.mean(np.linalg.norm(grads_right, axis=1))
        total_grad = avg_grad_left + avg_grad_right
        
        # 更新显示文本
        self.grad_para.text = (
            f"梯度范数: 左 {avg_grad_left:.2e} | "
            f"右 {avg_grad_right:.2e} | "
            f"总 {total_grad:.2e}"
        )

    def _update_gradient_display(self, attr, old, new):
        """梯度显示切换（双重清除保障）"""
        types = ["none", "geodesic", "geometry", "ht2t"]
        self.selected_gradient = types[new]
        
        # 立即执行清除操作
        self._plot_gradients(np.zeros((self.point_count,2)),
                        np.zeros((self.point_count,2)))
        
        if self.selected_gradient != "none":
            self._update_gradients()
    
    def _update_point_config(self, attr, old, new):
        """更新点配置（保持数据源独立性）"""
        if new == 0:  # 预设点
            self.point_count = 4  # 强制设置点数量为4
            # 正方形坐标（保持左右独立但初始相同）
            preset_points = np.array([[-1,-1], [-1,1], [1,1], [1,-1]])
            colors = generate_colors(4)
            
            # 完全替换左图数据源
            self.source_left.data = {
                'x': preset_points[:,0].tolist(),
                'y': preset_points[:,1].tolist(),
                'color': colors,
                'color_hex': colors
            }
            
            # 完全替换右图数据源
            self.source_right.data = {
                'x': preset_points[:,0].tolist(),
                'y': preset_points[:,1].tolist(),
                'color': colors,
                'color_hex': colors
            }
        else:  # 随机点
            self.point_count = 4  # 强制设置点数量为8
            colors = generate_colors(4)
            
            # 生成独立随机数据（左右初始相同但后续独立）
            new_data_left = {
                'x': (np.random.rand(self.point_count)*10).tolist(),
                'y': (np.random.rand(self.point_count)*10).tolist(),
                'color': colors,
                'color_hex': colors
            }
            new_data_right = {
                'x': (np.random.rand(self.point_count)*10).tolist(),
                'y': (np.random.rand(self.point_count)*10).tolist(),
                'color': colors,
                'color_hex': colors
            }
            
            # 完全替换数据源
            self.source_left.data = new_data_left
            self.source_right.data = new_data_right
        
        # 更新界面组件
        self._update_display()


    def _plot_gradients(self, grads_left, grads_right):
        """可视化梯度方向（最终解决方案）"""
        # ===== 增强清除阶段 =====
        print("\n[清除阶段]")
        for plot in [self.plot_left, self.plot_right]:
            # 精准识别需要清除的箭头
            to_remove = [
                r for r in plot.renderers 
                if isinstance(r, Arrow) and hasattr(r, 'tags') and "gradient_arrow" in r.tags
            ]
            
            # 创建新渲染器列表（排除梯度箭头）
            new_renderers = [r for r in plot.renderers if r not in to_remove]
            
            # 强制更新渲染器
            plot.renderers = new_renderers
            print(f"{plot.title.text} | 清除箭头: {len(to_remove)}，剩余渲染器: {len(new_renderers)}")

        # ===== 数据校验阶段 =====
        try:
            print("\n[数据校验]")
            assert len(self.source_left.data['x']) == self.point_count, "左图数据维度异常"
            assert len(self.source_right.data['x']) == self.point_count, "右图数据维度异常"
            assert grads_left.shape == (self.point_count, 2), f"左梯度形状异常: {grads_left.shape}"
            assert grads_right.shape == (self.point_count, 2), f"右梯度形状异常: {grads_right.shape}"
        except Exception as e:
            print(f"数据校验失败: {str(e)}")
            return

        # ===== 绘图阶段 =====
        if self.selected_gradient != "none":
            print("\n[绘制新箭头]")
            # 动态缩放计算
            max_grad = max(
                np.max(np.linalg.norm(grads_left, axis=1)),
                np.max(np.linalg.norm(grads_right, axis=1)),
                1e-6  # 防止除零
            )
            scale = 2.0 / max_grad
            
            # 生成唯一标签组
            arrow_tags = ["gradient_arrow", f"type_{self.selected_gradient}"]
            
            try:
                # 创建临时列表保存新箭头
                new_left_arrows = []
                new_right_arrows = []
                
                for i in range(self.point_count):
                    # 左图箭头
                    left_arrow = Arrow(
                        x_start=self.source_left.data['x'][i],
                        y_start=self.source_left.data['y'][i],
                        x_end=self.source_left.data['x'][i] + grads_left[i,0] * scale,
                        y_end=self.source_left.data['y'][i] + grads_left[i,1] * scale,
                        end=OpenHead(line_color="#E67E22", size=8),
                        line_color="#E67E22",
                        line_width=1.8,
                        tags=arrow_tags.copy()  # 必须使用copy创建新列表
                    )
                    new_left_arrows.append(left_arrow)
                    
                    # 右图箭头
                    right_arrow = Arrow(
                        x_start=self.source_right.data['x'][i],
                        y_start=self.source_right.data['y'][i],
                        x_end=self.source_right.data['x'][i] + grads_right[i,0] * scale,
                        y_end=self.source_right.data['y'][i] + grads_right[i,1] * scale,
                        end=OpenHead(line_color="#E67E22", size=8),
                        line_color="#E67E22",
                        line_width=1.8,
                        tags=arrow_tags.copy()
                    )
                    new_right_arrows.append(right_arrow)
                
                # 批量添加箭头（确保原子操作）
                self.plot_left.renderers.extend(new_left_arrows)
                self.plot_right.renderers.extend(new_right_arrows)
                
                # 强制触发渲染器更新
                self.plot_left.renderers = self.plot_left.renderers
                self.plot_right.renderers = self.plot_right.renderers
                
                print(f"成功创建 {len(new_left_arrows)*2} 个箭头")
                print(f"左图实际箭头数: {len(new_left_arrows)}")
                print(f"右图实际箭头数: {len(new_right_arrows)}")

            except Exception as e:
                print(f"箭头创建失败: {str(e)}")
        else:
            print("梯度显示已关闭")

    @property
    def layout(self):
        """最终布局（修复控件引用）"""
        return column(
            self.header,
            row(
                column(self.plot_left, self.editable_table_left, self.centroid_table_left),
                column(self.plot_right, self.editable_table_right, self.centroid_table_right),
                spacing=30
            ),
            self.diff_display,
            sizing_mode='fixed',
            margin=(20, 20, 20, 20)
        )

# === 应用启动 ===
app = VisualizationApp(k_neighbors=2)
curdoc().add_root(app.layout)
curdoc().title = "ht2t可视化（欧式版本）"